import os

from llama_index.agent.openai import OpenAIAgent
from llama_index.core.query_engine import SubQuestionQueryEngine
from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.embeddings.openai import OpenAIEmbedding
#pip install llama-hub unstructured

from llama_index.readers.file import UnstructuredReader
from pathlib import Path
from llama_index.core import VectorStoreIndex, StorageContext, load_index_from_storage, SimpleDirectoryReader
from llama_index.core import Settings

from llamaIndex.CommonClient import llm


# loader = UnstructuredReader()
# pdfData=loader.load_data(
#         file=Path(f"./data/nke-10k-2023.pdf"), split_documents=False
#     )
openAIEmbeddings=OpenAIEmbedding(api_key="sk-CftUbVSsA61lwwgMz9xvt6znTunQZfgBP8ZCVLbQsKfXUR6k",
    model='text-embedding-3-small',
    api_base="https://www.henapi.top/v1")
index_dir="./storage/pdf1"
if not os.path.exists(index_dir):
    pdfData = SimpleDirectoryReader("./data").load_data()


    Settings.chunk_size = 512
    storage_context = StorageContext.from_defaults()

    cur_index = VectorStoreIndex.from_documents(
            pdfData,
            storage_context=storage_context,
    embed_model=openAIEmbeddings,show_progress=True
    )
    storage_context.persist(persist_dir=f"./storage/pdf1")


storage_context = StorageContext.from_defaults(
            persist_dir=f"./storage/pdf1"
        )
index_set = load_index_from_storage(
            storage_context,
embed_model=openAIEmbeddings,show_progress=True
        )

individual_query_engine_tools = [
    QueryEngineTool(
        query_engine=index_set.as_query_engine(llm=llm),
        metadata=ToolMetadata(
            name=f"vector_index_pdf",
            description=f"useful for when you want to answer queries about the  SEC 10-K for nike",
        ),
    )
]

query_engine = SubQuestionQueryEngine.from_defaults(
    query_engine_tools=individual_query_engine_tools,
    llm=llm,
)

query_engine_tool = QueryEngineTool(
    query_engine=query_engine,
    metadata=ToolMetadata(
        name="sub_question_query_engine",
        description="useful for when you want to answer queries that require analyzing multiple nike",
    ),
)
tools = individual_query_engine_tools + [query_engine_tool]

agent = OpenAIAgent.from_tools(tools, verbose=True,llm=llm)
#普通问题
response = agent.chat("hi, i am bob")
print(str(response))
#关于索引数据相关的问题
response = agent.chat(
    "What was Nike's revenue in 2023?"
)
print(str(response))

cross_query_str = "Compare/contrast the risk factors described in the Uber 10-K across years. Give answer in bullet points."

response = agent.chat(cross_query_str)
print(str(response))
# 问的问题：What was Nike's revenue in 2023?
while True:
    text_input = input("User: ")
    if text_input == "exit":
        break
    response = agent.chat(text_input)
    print(f"Agent: {response}")